# 功能2、功能5
import warnings

warnings.filterwarnings('ignore')

import torch
import torch.nn as nn
import copy
import random
from tqdm import trange
from utils.options import args_parser
from utils.sampling import iid
from utils.dataset import load_data
from utils.test import test_img
from utils.attack import compromised_clients
from utils.update import BenignUpdate
import tenseal as ts
from pypbc import *
import numpy as np
import hashlib

Hash = hashlib.sha256
from utils.masked_input import Flatten_paremeter, KeyGen, PPGen, Hhash, sign_key_gen, sign, Verify_agg_sign_2, \
    decompose_agg_w_new, \
    Hhash_equa, Verify_agg_sign_1, Restore_values_to_w, Invert_Quality
from utils.trace import Tracing
import base64
from flask import Flask, request, jsonify, render_template
from flask_cors import CORS

app = Flask(__name__)
CORS(app)


server = None

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(784, 10)  # 784，100， 10
        self.fc2 = nn.Linear(10, 10)

    def forward(self, x):
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        return x


class ServerAgg:
    def __init__(self, args):
        self.args = args
        self.device = args.device
        self.method = args.method
        self.global_ep = args.global_ep
        self.num_clients = args.num_clients
        self.idxs_users = []
        self.compromised_idxs = []
        self.frac = args.frac
        # self.dict_users=[]
        self.net_glob = Net().to(args.device)
        # self.w_glob = []

        self.w_locals = []
        self.dataset_train=None
        # self.context = ''
        self.request_tracing_clients = [False]*self.num_clients
        self.num_request_tracing_clients = 0
        self.verify_result = []
        self.local_settings = [0]*self.num_clients

    def local_train(self,local_settings):
        args=self.args
        dataset_train = self.dataset_train
        # self.compromised_idxs=compromised_idxs
        print("self.device",self.device)
        print("local_settings = ", self.local_settings)

        for idx in self.idxs_users:
            # 生成签名的公私密钥对
            sk_sign_1_idx, pk_sign_1_idx = sign_key_gen(self.G2, self.g2, self.pairing)
            sk_sign_2_idx, pk_sign_2_idx = sign_key_gen(self.G2, self.g2, self.pairing)
            self.pk_sign_1_s[idx] = pk_sign_1_idx
            self.pk_sign_2_s[idx] = pk_sign_2_idx

            len_hv = 130 * args.num_clients
            hv_ini = [0] * len_hv

            print("执行loacl_train")
            # malicious client恶意客户端
            if local_settings[idx] != 0:
                # 获得本地参数
                local = BenignUpdate(args=args, dataset=dataset_train, idxs=self.dict_users[idx])
                w = local.train(net=copy.deepcopy(self.net_glob).to(args.device))
                # 参数铺平并量化
                w_flat = Flatten_paremeter(copy.deepcopy(w))
                # 生成线性同态哈希值
                if local_settings[idx]==1:
                    '''----------------------------伪造本地线性同态哈希值---------------------------------'''
                    w_flat_forged = [-num for num in w_flat]
                    h1_idx = Hhash(self.pp1, w_flat_forged, self.total_params, self.pairing)  # 同态哈希h1u
                    print(idx,"伪造哈希")
                else :
                    '''-------------正常的--------------'''
                    h1_idx = Hhash(self.pp1, w_flat, self.total_params, self.pairing)  # 同态哈希h1u

                h1_idx_str = Element.__str__(h1_idx)  # 哈希值转为字符串
                h1_idx_str_list = list(h1_idx_str)
                h1_idx_ASCII_list = [ord(mess) for mess in h1_idx_str_list]

                # 构建本地向量
                a = idx * 130
                b = a + 130
                replace_indices = list(range(a, b))
                for i, j in enumerate(replace_indices):
                    hv_ini[j] = h1_idx_ASCII_list[i]

                w_new_idx = hv_ini + w_flat

                # 同态加密 wu_new -> [[wu_new]]
                enc_w_new_idx = ts.bfv_vector(self.context, w_new_idx)
                # 序列化
                local_BFV_bytes = enc_w_new_idx.serialize()  # 序列化；bytes

                h2_idx_pre = Hhash(self.pp2, hv_ini, self.len_pp2, self.pairing)
                h2_idx = Element.__mul__(h1_idx, h2_idx_pre)  # h2u

                # 数字签名1
                if local_settings[idx] == 2 :
                    '''----------伪造第1个签名------------'''
                    sk_sign_1_forged, pk_sign_1_forged = sign_key_gen(self.G2, self.g2, self.pairing)
                    sig1_idx = sign(sk_sign_1_forged, self.pairing, h2_idx)
                    print(idx, "伪造签名1")
                else :
                    '''------------正常的------------'''
                    sig1_idx = sign(sk_sign_1_idx, self.pairing, h2_idx)

                # 数字签名2
                local_BFV_str = str(local_BFV_bytes)
                str_enc_w_new_sigm_idx = local_BFV_str + str(sig1_idx)
                h3_idx_str = Hash(str_enc_w_new_sigm_idx.encode()).hexdigest()
                h3_idx = Element.from_hash(self.pairing, self.G2, h3_idx_str)
                if local_settings[idx] == 3:
                    '''----------------伪造第2个签名--------------# 在Aggregation阶段的验证结果为0，并重新开始，未追踪恶意用户'''
                    sk_sign_2_forged, pk_sign_2_forged = sign_key_gen(self.G2, self.g2, self.pairing)
                    sig2_idx = sign(sk_sign_2_forged, self.pairing, h3_idx)
                    print(idx, "伪造签名2")
                else:
                    '''------------------正常的--------------------'''
                    sig2_idx = sign(sk_sign_2_idx, self.pairing, h3_idx)

                self.w_locals.append(copy.deepcopy(w_flat))

            # honest client可信客户端设置
            else:
                # local model parameter
                local = BenignUpdate(args=args, dataset=dataset_train, idxs=self.dict_users[idx])
                w = local.train(net=copy.deepcopy(self.net_glob).to(args.device))
                # 参数铺平并量化
                w_flat = Flatten_paremeter(copy.deepcopy(w))

                # 生成线性同态哈希值
                h1_idx = Hhash(self.pp1, w_flat, self.total_params, self.pairing)  # 同态哈希h1u

                h1_idx_str = Element.__str__(h1_idx)  # 哈希值转为字符串
                h1_idx_str_list = list(h1_idx_str)
                h1_idx_ASCII_list = [ord(mess) for mess in h1_idx_str_list]

                # 构建本地向量
                a = idx * 130
                b = a + 130
                replace_indices = list(range(a, b))
                for i, j in enumerate(replace_indices):
                    hv_ini[j] = h1_idx_ASCII_list[i]
                w_new_idx = hv_ini + w_flat

                # 同态加密 wu_new -> [[wu_new]]
                enc_w_new_idx = ts.bfv_vector(self.context, w_new_idx)
                # 序列化
                local_BFV_bytes = enc_w_new_idx.serialize()  # 序列化；bytes

                h2_idx_pre = Hhash(self.pp2, hv_ini, self.len_pp2, self.pairing)
                h2_idx = Element.__mul__(h1_idx, h2_idx_pre)  # h2u

                # 数字签名1
                sig1_idx = sign(sk_sign_1_idx, self.pairing, h2_idx)

                # 数字签名2
                local_BFV_str = str(local_BFV_bytes)
                str_enc_w_new_sigm_idx = local_BFV_str + str(sig1_idx)
                h3_idx_str = Hash(str_enc_w_new_sigm_idx.encode()).hexdigest()
                h3_idx = Element.from_hash(self.pairing, self.G2, h3_idx_str)

                sig2_idx = sign(sk_sign_2_idx, self.pairing, h3_idx)

                self.w_locals.append(copy.deepcopy(w_flat))

            # 存储信息
            # 保存本地向量密文
            str_idx = str(idx)
            filename0 = './local_BFV/local_BFV_' + str_idx + '.txt'
            with open(filename0, 'wb') as f:
                f.write(local_BFV_bytes)
            f.close()

            # 保存数字签名1
            str_sig1_idx = Element.__str__(sig1_idx)
            filename1 = './local_sig1/local_sig1_' + str_idx + '.txt'
            with open(filename1, 'w') as f:
                f.write(str_sig1_idx)
            f.close()

            # 保存数字签名2
            filename2 = './local_sig2/local_sig2_' + str_idx + '.txt'
            str_sig2_idx = Element.__str__(sig2_idx)
            with open(filename2, 'w') as f:
                f.write(str_sig2_idx)
            f.close()

        # self.global_func()


    def aggregation(self,server_state):
        args = self.args
        enc_w_new_s_recover, sig1_s, sig1_s_recover, sig2_s_recover = [], [], [], []
        if args.method == 'VeriTrac':
            # 读取数据
            for j in range(len(self.idxs_users)):
                str_idx = str(j)

                filename0 = './local_BFV/local_BFV_' + str_idx + '.txt'
                with open(filename0, 'rb') as f:
                    local_BFV_bytes_recover = f.read()
                f.close()
                enc_w_new_s_recover.append(local_BFV_bytes_recover)

                filename1 = './local_sig1/local_sig1_' + str_idx + '.txt'
                with open(filename1, 'r') as f1:
                    str_sig1_idx_recover = f1.read()
                f1.close()
                sig1_s_recover.append(str_sig1_idx_recover)

                filename2 = './local_sig2/local_sig2_' + str_idx + '.txt'
                with open(filename2, 'r') as f2:
                    str_sig2_idx_recover = f2.read()
                f2.close()
                sig2_idx_recover = Element(self.pairing, self.G2, value=str_sig2_idx_recover)
                sig2_s_recover.append(sig2_idx_recover)

            '''-------------Aggregation------------'''

            # 服务器验证客户端
            verify_aggregation = Verify_agg_sign_2(enc_w_new_s_recover, self.pairing, sig1_s_recover, self.pk_sign_2_s,
                                                   args.num_clients, self.g2, sig2_s_recover)
            print("verify_aggregation = ", verify_aggregation)

            # if verify_aggregation == 1 and -1 in self.compromised_idxs:
            if verify_aggregation == 1 and server_state == True:
                # if verify == 1 and -1 in compromised_idxs:
                # 聚合参数密文

                '''----------------------------伪造聚合模型---------------------------------'''
                print("server:伪造聚合模型")
                len_h = 130 * args.num_clients
                h_ini = [0] * len_h
                w_local_forged = h_ini + self.w_locals[0]
                enc_w_local_forged = ts.bfv_vector(self.context, w_local_forged)

                w_new_s_BFV = ts.bfv_vector_from(self.context, enc_w_new_s_recover[args.num_clients - 1])  # 反序列化
                print("----------------------------------------------------------------------------")
                for i in range(args.num_clients - 1):
                    w_new_s_BFV = w_new_s_BFV + ts.bfv_vector_from(self.context, enc_w_new_s_recover[i])
                w_new_s_BFV = w_new_s_BFV + enc_w_local_forged

                # 存储聚合结果
                filename4 = './agg_BFV/agg_BFV_w_new' + '.txt'
                local_BFV_bytes = w_new_s_BFV.serialize()  # 序列化；bytes
                with open(filename4, 'wb') as f:
                    f.write(local_BFV_bytes)
                f.close()

                '''----------------------------正常的---------------------------------'''
            # elif verify_aggregation == 1 and -1 not in self.compromised_idxs:
            elif verify_aggregation == 1 and server_state == False:
                w_new_s_BFV = ts.bfv_vector_from(self.context, enc_w_new_s_recover[args.num_clients - 1])  # 反序列化
                for i in range(args.num_clients - 1):
                    w_new_s_BFV = w_new_s_BFV + ts.bfv_vector_from(self.context, enc_w_new_s_recover[i])
                # 存储聚合结果
                filename4 = './agg_BFV/agg_BFV_w_new' + '.txt'
                local_BFV_bytes = w_new_s_BFV.serialize()  # 序列化；bytes
                with open(filename4, 'wb') as f:
                    f.write(local_BFV_bytes)
                f.close()
        return verify_aggregation


    def client_verify(self):
        args=self.args
        pk_sign_1_s=self.pk_sign_1_s
        '''unmasking and verification'''
        idxs_users = self.idxs_users
        for idx in idxs_users:
            # if idx not in self.compromised_idxs:
                # 读取数据
                filename4 = './agg_BFV/agg_BFV_w_new' + '.txt'
                with open(filename4, 'rb') as f:
                    agg_BFV_bytes_recover = f.read()
                f.close()

                # 解密
                enc_agg_bfv_recover = ts.bfv_vector_from(self.context, agg_BFV_bytes_recover)  # 反序列化
                decrypted_agg_w_new = enc_agg_bfv_recover.decrypt()

                # 分解向量
                decrypted_agg_w, h1_idx_list_recover = decompose_agg_w_new(decrypted_agg_w_new, idxs_users,
                                                                           self.len_pp2,
                                                                           args.num_clients)

                # 判断是否满足同态哈希的同态性
                # 求同态哈希值
                verify1 = Hhash_equa(self.pp1, decrypted_agg_w, len(decrypted_agg_w), self.pairing, self.G2,
                                     h1_idx_list_recover)
                print(idx,"verify1 = ", verify1)

                if verify1 == 0:
                    # self.num_request_tracing_clients += 1

                    filename_1 = './Verify/user_verify_result' + '.txt'
                    # local_BFV_bytes = w_new_s_BFV.serialize()  # 序列化；bytes
                    self.verify_result[idx] = False
                    with open(filename_1, 'w') as f:
                        f.write("The verification failed.")
                    f.close()

                elif verify1 == 1:
                    # 利用数字签名1验证各个哈希值的真实性
                    # 读取数字签名1
                    sig1_s_ele_recover = []
                    for i in range(len(idxs_users)):
                        str_idx = str(idxs_users[i])
                        filename1 = './local_sig1/local_sig1_' + str_idx + '.txt'
                        with open(filename1, 'r') as f1:
                            str_sig1_idx_recover = f1.read()
                        f1.close()
                        ele_sig1_idx_recover = Element(self.pairing, self.G2, value=str_sig1_idx_recover)
                        sig1_s_ele_recover.append(ele_sig1_idx_recover)

                    h2_idx_rec_list = []
                    for i, h1_idx_recover in zip(range(len(idxs_users)), h1_idx_list_recover):
                        len_hv = 130 * args.num_clients
                        hv_ini_local = [0] * len_hv

                        # 对扩展向量求同态哈希
                        h1_idx_recover_list = list(h1_idx_recover)
                        h1_idx_recover_ASCII = [ord(mess) for mess in h1_idx_recover_list]

                        a = i * 130
                        b = a + 130
                        replace_indices = list(range(a, b))
                        for ii, jj in enumerate(replace_indices):
                            hv_ini_local[jj] = h1_idx_recover_ASCII[ii]

                        h2_idx_pre = Hhash(self.pp2, hv_ini_local, self.len_pp2, self.pairing)
                        h1_idx_ele_rec = Element(self.pairing, self.G2, value=h1_idx_recover)
                        h2_idx_rec = Element.__mul__(h1_idx_ele_rec, h2_idx_pre)  # h2u
                        h2_idx_rec_list.append(h2_idx_rec)

                    verify2 = Verify_agg_sign_1(self.pairing, h2_idx_rec_list, sig1_s_ele_recover, pk_sign_1_s,
                                                args.num_clients, self.g2)
                    print("verify2 = ", verify2)

                    if verify2 == 0:
                        # self.num_request_tracing_clients += 1
                        filename_2 = './Verify/user_verify_result' + '.txt'
                        self.verify_result[idx] = False
                        # local_BFV_bytes = w_new_s_BFV.serialize()  # 序列化；bytes
                        with open(filename_2, 'w') as f:
                            f.write("The verification failed.")
                        f.close()
                    elif verify2 == 1:
                        aggregated_parameter = Invert_Quality(decrypted_agg_w, args.num_clients)
                        self.w_glob = Restore_values_to_w(copy.deepcopy(self.w_glob), decrypted_agg_w)

    def trace(self):
        args = self.args
        idxs_users=self.idxs_users
        if self.num_request_tracing_clients > args.num_clients / 2:
            unhonest_parts = Tracing(args, idxs_users, self.pk_sign_1_s, self.pk_sign_2_s, self.pairing, self.G2, self.g2,
                                            self.context,
                                            self.len_pp2,
                                            args.num_clients, self.pp2, self.pp1)
            print("unhonest_parts = ", unhonest_parts)

            filename = './TRACE/tracing_result' + '.txt'
            # local_BFV_bytes = w_new_s_BFV.serialize()  # 序列化；bytes
            with open(filename, 'w') as f:
                f.write(str(unhonest_parts))
            f.close()

        return unhonest_parts


    def set_clients(self):
        self.net_glob.load_state_dict(self.w_glob)


    def global_func(self):
        args = self.args

        dataset_test = self.dataset_test



        # 追踪服务器追踪
        self.trace(args)

        self.net_glob.load_state_dict(self.w_glob)

        test_acc, test_loss = test_img(self.net_glob.to(args.device), dataset_test, args)

        print(f"Round: {iter}")
        print(f"Test accuracy: {test_acc}")
        print(f"Test loss: {test_loss}")

    def func_main(self, args):
        random.seed(args.seed)
        np.random.seed(args.seed)
        torch.manual_seed(args.seed)
        torch.cuda.manual_seed_all(args.seed)

        self.dataset_train, self.dataset_test, dataset_val = load_data(args)

        # sample users
        self.dict_users = iid(self.dataset_train, args)

        self.net_glob.train()

        # copy weights
        self.w_glob = self.net_glob.state_dict()
        # 设置恶意参与方（服务器、客户端）
        # compromised_idxs = [-1,1,2]
        # self.compromised_idxs = compromised_idxs
        # print("compromised_idxs = ", self.compromised_idxs)

        '''初始化'''
        self.context = ts.context(ts.SCHEME_TYPE.BFV, poly_modulus_degree=4096, plain_modulus=1032193)
        [self.G2, self.g2, self.params, self.pairing] = KeyGen(512, 160)
        with torch.no_grad():
            self.total_params = len(Flatten_paremeter(copy.deepcopy(self.w_glob)))

        self.pp1 = PPGen(self.g2, self.params, self.pairing, self.total_params)  # 参数长度 d
        self.len_pp2 = args.num_clients * 130
        self.pp2 = PPGen(self.g2, self.params, self.pairing, self.len_pp2)

        # for iter in trange(args.global_ep):
        #     self.global_func( args, dataset_train, dataset_test, total_params)
        args = self.args

        # for iter in trange(args.global_ep):

        selected_clients = max(int(args.frac * args.num_clients), 1)
        self.idxs_users = np.random.choice(range(args.num_clients), selected_clients, replace=False)

        ''' Masked Input '''
        # print("idxs_users = ", self.idxs_users)
        self.pk_sign_1_s = [0 for i in range(args.num_clients)]
        self.pk_sign_2_s = [0 for i in range(args.num_clients)]



@app.route("/run", methods=['GET', 'POST'])
def hello1():
    # parse args
    global server
    args = args_parser()
    # args.device = torch.device('cuda:{}'.format(args.gpu) if torch.cuda.is_available() and args.gpu != -1 else 'cpu')
    args.device = torch.device('cpu')
    args.method = 'VeriTrac'
    args.global_ep = 1
    args.num_clients = 10
    server = ServerAgg(args)
    data = request.json
    # 假设前端传递一个 client_ids 数组，用于指定要设置的客户端
    # compromised_idxs = data.get('compromised_idxs', [])

    server.func_main(args)
    return 'Hello, World!'

@app.route("/global_func", methods=['GET', 'POST'])
def global_function():
    global server

    server.global_func()
    return 'Global_func done!'

@app.route("/local_train", methods=['GET', 'POST'])
def local_train_function():
    global server
    data = request.json
    # 假设前端传递一个 client_ids 数组，用于指定要设置的客户端
    server.local_settings = data.get('local_settings', [])
    server.local_train(server.local_settings)
    return 'local_train done!'

@app.route("/post_model", methods=['GET', 'POST'])
def post_model():
    global server

    return "post_model done!"

@app.route("/server_receive_model", methods=['GET', 'POST'])
def server_receive_model():
    global server

    enc_w_new_s_recover = []
    for j in range(len(server.idxs_users)):
        str_idx = str(j)
        filename0 = './local_BFV/local_BFV_' + str_idx + '.txt'
        try:
            with open(filename0, 'rb') as f:
                local_BFV_bytes_recover = f.read()
            enc_w_new_s_recover.append(local_BFV_bytes_recover)
        except FileNotFoundError:
            return jsonify({"error": f"File {filename0} not found"}), 404

    # 对每个字节数据进行 base64 编码
    encoded_data = [base64.b64encode(data).decode('utf-8') for data in enc_w_new_s_recover]

    response_data = {"enc_w_new_s_recover": encoded_data}
    # server.global_func()
    return jsonify(response_data)

@app.route("/aggregate", methods=['GET', 'POST'])
def aggregate():
    global server
    data = request.json
    server_state = data.get('server_state')
    # 全局聚合
    verify_aggregation = server.aggregation(server_state)
    agg_model = []
    filename = './agg_BFV/agg_BFV_w_new' + '.txt'
    try:
        with open(filename, 'rb') as f:
            local_BFV_bytes_recover = f.read()
        agg_model.append(local_BFV_bytes_recover)
    except FileNotFoundError:
        return jsonify({"error": f"File {filename} not found"}), 404
    # 对每个字节数据进行 base64 编码
    encoded_agg_model = [base64.b64encode(data).decode('utf-8') for data in agg_model]

    return jsonify({'verify_aggregation':verify_aggregation,'aggregation_model':encoded_agg_model})

@app.route("/set_clients", methods=['GET', 'POST'])
def set_clients():
    global server
    return "set_clients done!"

@app.route("/client_receive_model", methods=['GET', 'POST'])
def client_receive_model():
    global server
    return "client_receive_model done!"

@app.route("/client_verify", methods=['GET', 'POST'])
def client_verify():
    global server
    # 客户端验证、可选是否验证
    server.num_request_tracing_clients = 0
    server.verify_result=[True] * len(server.idxs_users)
    server.client_verify()

    return jsonify(server.verify_result)

@app.route("/trace_request", methods=['GET', 'POST'])
def trace_request():
    global server
    data = request.json

    request_id = data.get('idx' )
    server.request_tracing_clients[request_id] = True

    return jsonify({'request_id': request_id})

@app.route("/start_trace", methods=['GET', 'POST'])
def start_trace():
    global server
    # data = request.json
    # request_id = data.get('idx' )

    server.num_request_tracing_clients = server.request_tracing_clients.count(True)
    unhonest_participants = []
    for char in server.trace():
        if char.isdigit():
            unhonest_participants.append(int(char))
        else:
            unhonest_participants.append(-1)
    # unhonest_participants = [int(char) for char in server.trace()]
    return jsonify({'unhonest_participants': unhonest_participants})

@app.route("/return_trace_result", methods=['GET', 'POST'])
def return_trace_result():
    global server
    return "return_trace_result done!"


if __name__ == '__main__':
    # 你的本机 IP 地址，可以在命令行中使用 ipconfig（Windows）或 ifconfig（Linux/Mac）查找
    host = '127.0.0.1'  # 替换成你的本机 IP 地址
    port = 50000  # 你想要使用的端口号
    app.run(host=host, port=port)
